{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "df0dfd71",
   "metadata": {
    "colab_type": "text",
    "id": "OE4xjKWgtIX2"
   },
   "source": [
    "# Bunny ICP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83574180",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KdXAerwV13rQ",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "# Imports\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    PointCloud,\n",
    "    Rgba,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    StartMeshcat,\n",
    ")\n",
    "from scipy.spatial import KDTree\n",
    "\n",
    "from manipulation import FindResource\n",
    "from manipulation.exercises.grader import Grader\n",
    "from manipulation.exercises.pose.test_icp import TestICP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ce6fcf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f9b6150",
   "metadata": {
    "colab_type": "text",
    "id": "jigwRNW7tIYQ"
   },
   "source": [
    "## Problem Description\n",
    "In the lecture, we learned about the Iterative Closest Point (ICP) algorithm. In this exercise, you will implement the ICP algorithm to solve the standard Stanford Bunny problem!\n",
    "\n",
    "**These are the main steps of the exercise:**\n",
    "1. Implement the ```least_squares_transform``` function to optimize transformation given correspondence\n",
    "2. Implement the ```icp``` algorithm using the functions implemented above.\n",
    "\n",
    "Let's first visualize the point clouds of Stanford bunny in meshcat!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79be7c41",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "VnQvdI6nOl4d",
    "outputId": "866d6e69-d188-4669-8b01-825d8e616b0d"
   },
   "outputs": [],
   "source": [
    "# Visualize Stanford Bunny\n",
    "xyzs = np.load(FindResource(\"models/bunny/bunny.npy\"))\n",
    "cloud = PointCloud(xyzs.shape[1])\n",
    "cloud.mutable_xyzs()[:] = xyzs\n",
    "\n",
    "# Pose for the blue bunny\n",
    "X_blue = RigidTransform(\n",
    "    RotationMatrix.MakeXRotation(np.pi / 6), [-0.1, 0.1, 0.1]\n",
    ")\n",
    "\n",
    "pointcloud_model = xyzs\n",
    "pointcloud_scene = X_blue.multiply(xyzs)\n",
    "\n",
    "meshcat.Delete()\n",
    "meshcat.SetProperty(\"/Background\", \"visible\", False)\n",
    "meshcat.SetProperty(\"/Cameras/default/rotated/<object>\", \"zoom\", 10.5)\n",
    "meshcat.SetObject(\"red_bunny\", cloud, point_size=0.01, rgba=Rgba(1.0, 0, 0))\n",
    "meshcat.SetTransform(\"red_bunny\", RigidTransform())\n",
    "meshcat.SetObject(\"blue_bunny\", cloud, point_size=0.01, rgba=Rgba(0, 0, 1.0))\n",
    "meshcat.SetTransform(\"blue_bunny\", X_blue)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c007f1e",
   "metadata": {
    "colab_type": "text",
    "id": "abo92_2stIYW"
   },
   "source": [
    "## Point cloud registration with known correspondences\n",
    "\n",
    "In this section, you will follow the [derivation](http://manipulation.csail.mit.edu/pose.html#registration) to solve the optimization problem below.\n",
    "\n",
    "$$\\begin{aligned} \\min_{p\\in\\mathbb{R}^3,R\\in\\mathbb{R}^{3\\times3}} \\quad & \\sum_{i=1}^{N_s} \\| p + R \\hspace{.1cm} {^Op^{m_{c_i}}} - p^{s_i}\\|^2, \\\\ s.t. \\quad & RR^T = I, \\quad \\det(R)=1\\end{aligned}$$\n",
    "\n",
    "The goal is to find the transform that registers the point clouds of the model and the scene, assuming the correspondence is known.  You may refer to the implementation from [deepnote](https://deepnote.com/workspace/Manipulation-ac8201a1-470a-4c77-afd0-2cc45bc229ff/project/4-Geometric-Pose-Estimation-cc6340f5-374e-449a-a195-839a3cedec4a/%2Ficp.ipynb) and the explanation from [textbook](http://manipulation.csail.mit.edu/pose.html#icp).\n",
    "\n",
    "In the cell below, implement the ```least_squares_transform``` nethod."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7279d213",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ll_FlqVotIYX"
   },
   "outputs": [],
   "source": [
    "def least_squares_transform(scene, model) -> RigidTransform:\n",
    "    \"\"\"\n",
    "    Calculates the least-squares best-fit transform that maps corresponding\n",
    "    points scene to model.\n",
    "    Args:\n",
    "      scene: 3xN numpy array of corresponding points\n",
    "      model: 3xM numpy array of corresponding points\n",
    "    Returns:\n",
    "      X_BA: A RigidTransform object that maps point_cloud_A on to point_cloud_B\n",
    "            such that\n",
    "                        X_BA.multiply(model) ~= scene,\n",
    "    \"\"\"\n",
    "    X_BA = RigidTransform()\n",
    "    ##################\n",
    "    # your code here\n",
    "    ##################\n",
    "    return X_BA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4d5bc0",
   "metadata": {
    "colab_type": "text",
    "id": "IejlqJ3vtIYg"
   },
   "source": [
    "## Point correspondence from closest point\n",
    "The ```least_squares_transform``` function assumes that the point correspondence is known. Unfortunately, this is often not the case, so we will have to estimate the point correspondence as well. A common heuristics for estimating point correspondence is the closest point/nearest neighbor.\n",
    "\n",
    "We have implemented the closest neighbors using [scipy's implementation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.html) of [k-d trees](https://en.wikipedia.org/wiki/K-d_tree)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a68ebd7",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_-bGj1a1OkbU"
   },
   "outputs": [],
   "source": [
    "def nearest_neighbors(scene, model):\n",
    "    \"\"\"\n",
    "    Find the nearest (Euclidean) neighbor in model for each\n",
    "    point in scene\n",
    "    Args:\n",
    "        scene: 3xN numpy array of points\n",
    "        model: 3xM numpy array of points\n",
    "    Returns:\n",
    "        distances: (N, ) numpy array of Euclidean distances from each point in\n",
    "            scene to its nearest neighbor in model.\n",
    "        indices: (N, ) numpy array of the indices in model of each\n",
    "            scene point's nearest neighbor - these are the c_i's\n",
    "    \"\"\"\n",
    "    distances = np.empty(scene.shape[1], dtype=float)\n",
    "    indices = np.empty(scene.shape[1], dtype=int)\n",
    "\n",
    "    kdtree = KDTree(model.T)\n",
    "    for i in range(model.shape[1]):\n",
    "        distances[i], indices[i] = kdtree.query(scene[:, i], 1)\n",
    "\n",
    "    return distances, indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26bc5cab",
   "metadata": {
    "colab_type": "text",
    "id": "KtvN0kBntIYo"
   },
   "source": [
    "## Iterative Closest Point (ICP)\n",
    "Now you should be able to register two point clouds iteratively by first finding/updating the estimate of point correspondence with ```nearest_neighbors``` and then computing the transform using ```least_squares_transform```. You may refer to the explanation from [textbook](http://manipulation.csail.mit.edu/pose.html#icp).\n",
    "\n",
    "**In the cell below, complete the implementation of ICP algorithm using the  ```nearest_neighbors``` and ```least_squares_transform``` methods from above.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27e5fa4f",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wETDMjk4tIYp"
   },
   "outputs": [],
   "source": [
    "def icp(scene, model, max_iterations=20, tolerance=1e-3):\n",
    "    \"\"\"\n",
    "    Perform ICP to return the correct relative transform between two set of points.\n",
    "    Args:\n",
    "        scene: 3xN numpy array of points\n",
    "        model: 3xM numpy array of points\n",
    "        max_iterations: max amount of iterations the algorithm can perform.\n",
    "        tolerance: tolerance before the algorithm converges.\n",
    "    Returns:\n",
    "      X_BA: A RigidTransform object that maps point_cloud_A on to point_cloud_B\n",
    "            such that\n",
    "                        X_BA.multiply(model) ~= scene,\n",
    "      mean_error: Mean of all pairwise distances.\n",
    "      num_iters: Number of iterations it took the ICP to converge.\n",
    "    \"\"\"\n",
    "    X_BA = RigidTransform()\n",
    "\n",
    "    mean_error = 0\n",
    "    num_iters = 0\n",
    "    prev_error = 0\n",
    "\n",
    "    while True:\n",
    "        num_iters += 1\n",
    "\n",
    "        # your code here\n",
    "        ##################\n",
    "\n",
    "        mean_error = np.inf  # Modify to add mean error.\n",
    "        ##################\n",
    "\n",
    "        if (\n",
    "            abs(mean_error - prev_error) < tolerance\n",
    "            or num_iters >= max_iterations\n",
    "        ):\n",
    "            break\n",
    "\n",
    "        prev_error = mean_error\n",
    "\n",
    "        meshcat.SetTransform(\"red_bunny\", X_BA)\n",
    "\n",
    "    return X_BA, mean_error, num_iters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64180ba2",
   "metadata": {
    "colab_type": "text",
    "id": "WChfoIVWtIYy"
   },
   "source": [
    "Now you should be able to visualize the registration of the Stanford bunny! Have fun!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91ab6e74",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "P8oIuMsCDMJM"
   },
   "outputs": [],
   "source": [
    "icp(pointcloud_scene, pointcloud_model, max_iterations=30, tolerance=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b517943",
   "metadata": {
    "colab_type": "text",
    "id": "ucRnypactIY2"
   },
   "source": [
    "## How will this notebook be Graded?\n",
    "\n",
    "If you are enrolled in the class, this notebook will be graded using [Gradescope](www.gradescope.com). You should have gotten the enrollement code on our announcement in Piazza.\n",
    "\n",
    "For submission of this assignment, you must:\n",
    "- Download and submit the notebook `bunny_icp.ipynb` to Gradescope's notebook submission section, along with your notebook for the other problems.\n",
    "\n",
    "We will evaluate the local functions in the notebook to see if the function behaves as we have expected. For this exercise, the rubric is as follows:\n",
    "- [3 pts] `least_squares_transform` must be implemented correctly.\n",
    "- [3 pts] `icp` must be implemented correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6037e9d",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "lfmnYSMItIY3",
    "outputId": "21d3a0f9-0e63-4eb7-cc7f-d95677e38c5e"
   },
   "outputs": [],
   "source": [
    "Grader.grade_output([TestICP], [locals()], \"results.json\")\n",
    "Grader.print_test_results(\"results.json\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}